# -*- coding: utf-8 -*-
# @File    : 4_muilt_get_username.py
# 多Cookie并发爬取用户信息

from Euclidweibo.Get_user_info import Get_user_info  # 获取基本信息（包含准确粉丝数）
from Euclidweibo.Get_user_info_v2 import Get_user_info_v2  # 获取详细信息（包含生日等）
import time
import datetime
import os
import json
import pandas as pd
import random
import threading
import concurrent.futures
from cookies_manager import CookiesManager
import re
from queue import Queue
import traceback
import requests


# 继承原有功能中的辅助函数
def determine_verified_type(verified_url):
    """
    根据verified_url判断认证类型
    
    Args:
        verified_url (str): 认证URL
    
    Returns:
        str: 认证类型描述
    """
    if not verified_url:
        return '未认证'
    
    # 商业认证
    if 'bluev' in verified_url or 'biz.weibo.com' in verified_url:
        return '企业认证'
    
    # 个人认证
    if 'verified.weibo.com' in verified_url:
        return '个人认证'
    
    # 其他类型认证
    return '已认证'


def determine_real_auth_status(user_info):
    """
    综合判断用户的实名认证状态
    
    Args:
        user_info (dict): 用户信息数据
    
    Returns:
        str: 实名认证状态描述
    """
    # 如果明确标记为已实名认证
    if user_info.get('real_auth') is True:
        return '是'
    
    # 如果有认证URL，通常也意味着实名认证
    if user_info.get('verified_url', ''):
        return '是'
        
    # 检查实名信息是否存在
    real_name = user_info.get('real_name', {})
    if real_name and (real_name.get('name') or real_name.get('career')):
        return '是'
    
    # 默认情况下返回否
    return '否'


def extract_followers_count_from_all_sources(user_info):
    """
    综合多种方法提取粉丝数
    
    Args:
        user_info (dict): 用户信息数据
    
    Returns:
        str: 粉丝数量
    """
    import re
    
    # 首先尝试从标签中找粉丝数，最可能包含真实数据
    if user_info.get('label_desc'):
        for label in user_info.get('label_desc', []):
            label_name = label.get('name', '')
            
            # 尝试匹配"粉丝数xx万+"的格式
            fans_match = re.search(r'粉丝数(\d+(\.\d+)?)万\+', label_name)
            if fans_match:
                fans_count = float(fans_match.group(1)) * 10000
                return str(int(fans_count))
                
            # 尝试匹配"xx万粉丝"格式
            fans_match = re.search(r'(\d+(\.\d+)?)\s*万粉丝', label_name)
            if fans_match:
                fans_count = float(fans_match.group(1)) * 10000
                return str(int(fans_count))
                
            # 尝试匹配"粉丝xx"格式
            fans_match = re.search(r'粉丝\s*(\d+)', label_name)
            if fans_match:
                return fans_match.group(1)
                
            # 尝试匹配"xx粉丝"格式(不带万字)
            fans_match = re.search(r'(\d+)\s*粉丝', label_name)
            if fans_match:
                return fans_match.group(1)
    
    # 从描述文本中可能包含的粉丝信息提取
    desc_text = user_info.get('desc_text', '')
    if desc_text:
        # 匹配描述中可能出现的粉丝数
        fans_match = re.search(r'(\d+(\.\d+)?)\s*万粉丝', desc_text)
        if fans_match:
            fans_count = float(fans_match.group(1)) * 10000
            return str(int(fans_count))
    
    # 尝试从个人描述中提取
    description = user_info.get('description', '')
    if description:
        # 匹配描述中可能出现的粉丝数
        fans_match = re.search(r'(\d+(\.\d+)?)\s*万粉丝', description)
        if fans_match:
            fans_count = float(fans_match.group(1)) * 10000
            return str(int(fans_count))
    
    # 最后才从API返回的followers字段获取，因为可能为0
    followers = user_info.get('followers', {})
    if followers and isinstance(followers, dict) and 'total_number' in followers:
        return str(followers.get('total_number', ''))
    
    return ''


def extract_friends_count(user_info):
    """
    提取好友数量
    
    Args:
        user_info (dict): 用户信息数据
    
    Returns:
        str: 好友数
    """
    friend_info = user_info.get('friend_info', '')
    
    # 处理字符串中的<a>标签
    if isinstance(friend_info, str) and friend_info:
        import re
        # 处理"他有 <a>xxx</a> 个好友"或"她有 <a>xxx</a> 个好友"的格式
        match = re.search(r'[他她]有\s*<a>(\d+)</a>', friend_info)
        if match:
            return match.group(1)
        
        # 处理"他有 <a>xxx"的格式(缺少结束标签)
        match = re.search(r'[他她]有\s*<a>(\d+)', friend_info)
        if match:
            return match.group(1)
    
    # 检查是否有别的位置存储好友数
    # 有些接口可能会直接返回数字
    if isinstance(friend_info, int) and friend_info > 0:
        return str(friend_info)
    
    return ''


def load_progress():
    """
    加载断点续爬进度
    
    Returns:
        set: 已处理过的用户ID集合
    """
    progress_file = 'Weibo/爬取进度/user_crawl_progress.json'
    if os.path.exists(progress_file):
        with open(progress_file, 'r', encoding='utf-8') as f:
            return set(json.load(f))
    return set()


def save_progress(processed_ids):
    """
    保存爬取进度
    
    Args:
        processed_ids (set): 已处理的用户ID集合
    """
    progress_file = 'Weibo/爬取进度/user_crawl_progress.json'
    # 确保目录存在
    os.makedirs('Weibo/爬取进度', exist_ok=True)
    with open(progress_file, 'w', encoding='utf-8') as f:
        json.dump(list(processed_ids), f)


def check_cookie_validity(cookie):
    """
    检查cookie是否有效
    
    Args:
        cookie (str): cookie字符串
    
    Returns:
        bool: cookie是否有效
    """
    test_url = "https://weibo.com/ajax/profile/info?uid=1669879400"  # 测试用的微博UID，用任意有效UID
    headers = {
        'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/104.0.0.0 Safari/537.36',
        'accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.9',
        'accept-encoding': 'gzip, deflate, br',
        'cookie': cookie,
        'accept-language': 'zh-CN,zh;q=0.9,en;q=0.8',
        'connection': 'keep-alive'
    }
    
    try:
        response = requests.get(test_url, headers=headers, timeout=10)
        
        # 检查HTTP状态码
        if response.status_code != 200:
            return False
        
        # 尝试解析响应内容
        response_json = response.json()
        
        # 检查响应中是否有错误码，表示登录失效
        if 'ok' in response_json and response_json['ok'] == 0:
            return False
        
        # 检查是否有用户数据
        if 'data' not in response_json:
            return False
            
        return True
    except Exception as e:
        print(f"检查cookie有效性时出错: {str(e)}")
        return False


class MultiThreadCrawler:
    """
    多线程爬虫类，使用多个Cookie并发爬取用户信息
    """
    
    def __init__(self, input_csv='Weibo/1_社交媒体特定表达_帖子内容.csv', output_csv='Weibo/2_社交媒体特定表达_用户信息内容.csv'):
        """
        初始化多线程爬虫
        
        Args:
            input_csv: 输入CSV文件路径，包含用户ID
            output_csv: 输出CSV文件路径，保存用户信息
        """
        self.input_csv = input_csv
        self.output_csv = output_csv
        self.cookies_manager = CookiesManager()
        self.processed_ids = set()  # 已处理的用户ID集合
        self.failed_ids = []  # 失败的用户ID列表
        self.results = []  # 结果列表
        self.results_lock = threading.Lock()  # 结果列表锁
        self.progress_lock = threading.Lock()  # 进度锁
        self.fail_lock = threading.Lock()  # 失败ID锁
        self.threads = []  # 线程列表
        self.progress_dir = 'Weibo/爬取进度'
        self.failed_ids_file = os.path.join(self.progress_dir, 'failed_user_ids.json')
        
        # 确保目录存在
        os.makedirs(self.progress_dir, exist_ok=True)
        
        # 加载现有数据
        self.initial_setup()
    
    def initial_setup(self):
        """
        初始化设置，加载已处理的用户ID和现有结果
        """
        # 加载已处理的用户ID
        self.processed_ids = load_progress()
        print(f"已加载 {len(self.processed_ids)} 个已处理的用户ID")
        
        # 加载现有结果
        if os.path.exists(self.output_csv):
            try:
                existing_df = pd.read_csv(self.output_csv)
                self.results = existing_df.to_dict('records')
                print(f"已加载 {len(self.results)} 条现有用户信息")
            except Exception as e:
                print(f"加载现有用户信息失败: {str(e)}")
                self.results = []
        
        # 加载失败的用户ID
        if os.path.exists(self.failed_ids_file):
            try:
                with open(self.failed_ids_file, 'r', encoding='utf-8') as f:
                    self.failed_ids = json.load(f)
                print(f"已加载 {len(self.failed_ids)} 个失败用户ID")
            except Exception as e:
                print(f"加载失败用户ID失败: {str(e)}")
                self.failed_ids = []
    
    def save_failed_ids(self):
        """
        保存失败的用户ID到文件
        """
        with self.fail_lock:
            with open(self.failed_ids_file, 'w', encoding='utf-8') as f:
                json.dump(self.failed_ids, f)
    
    def save_results(self):
        """
        保存结果到CSV文件
        """
        with self.results_lock:
            # 创建DataFrame并保存结果
            df = pd.DataFrame(self.results)
            df.to_csv(self.output_csv, index=False, encoding='utf-8-sig')
            print(f"已保存 {len(self.results)} 条用户信息到 {self.output_csv}")
    
    def split_user_ids(self):
        """
        按照Cookie数量拆分用户ID，分配给不同的Cookie处理
        
        Returns:
            Dict[str, List]: 每个Cookie ID对应的用户ID列表
        """
        # 读取CSV文件，获取所有用户ID
        df = pd.read_csv(self.input_csv)
        all_user_ids = df['uid'].unique().tolist()
        
        # 过滤出未处理的用户ID
        remaining_ids = [str(uid) for uid in all_user_ids if str(uid) not in self.processed_ids]
        print(f"总用户数: {len(all_user_ids)}")
        print(f"未处理用户数: {len(remaining_ids)}")
        
        # 如果没有未处理的用户ID，直接返回
        if not remaining_ids:
            print("所有用户都已处理完成！")
            return {}
        
        # 按照Cookie数量拆分用户ID
        # 确保multi_user_cookies目录存在
        os.makedirs(os.path.join(self.progress_dir, 'multi_user_cookies'), exist_ok=True)
        # 修改输出目录参数为multi_user_cookies
        return self.cookies_manager.split_tasks(remaining_ids, os.path.join(self.progress_dir, 'multi_user_cookies'))
    
    def process_user(self, uid, cookie_id):
        """
        处理单个用户ID，获取用户信息
        
        Args:
            uid: 用户ID
            cookie_id: 处理该用户的Cookie ID
            
        Returns:
            dict: 处理结果，包括用户信息和状态
        """
        try:
            # 获取该Cookie对应的header
            cookie_content, header = self.cookies_manager.get_cookie_with_header(cookie_id)
            
            # 检查cookie是否有效
            if not cookie_content:
                return {'status': 'error', 'uid': uid, 'message': f"Cookie {cookie_id} 不可用"}
            
            # 进一步检查cookie有效性
            if not check_cookie_validity(cookie_content):
                error_msg = f"Cookie {cookie_id} 已失效，请更新"
                print(f"\n[错误] {error_msg}")
                return {'status': 'error', 'uid': uid, 'message': error_msg}
            
            # 修改为0.5-1秒的随机延时
            time.sleep(random.uniform(0.5, 0.7))
            
            # 1. 首先获取详细信息(生日、简介等)
            try:
                detail_info = Get_user_info_v2(uid, header=header)
            except Exception as e:
                error_msg = f"获取用户详细信息失败: {str(e)}"
                print(f"\n[错误] {error_msg}")
                return {'status': 'error', 'uid': uid, 'message': error_msg}
            
            # 两次请求间增加延时
            time.sleep(random.uniform(0.5, 0.7))
            
            # 2. 获取基本信息(粉丝数等)
            try:
                basic_info = Get_user_info(uid, header=header)
            except Exception as e:
                error_msg = f"获取用户基本信息失败: {str(e)}"
                print(f"\n[错误] {error_msg}")
                return {'status': 'error', 'uid': uid, 'message': error_msg}
            
            # 3. 合并两个API的结果
            # 创建一个新的字典来存储合并后的数据
            user_info = detail_info.copy()
            
            # 如果详细信息中没有user字段，但基本信息中有，则添加
            if 'user' not in user_info and 'user' in basic_info:
                user_info['user'] = basic_info['user']
            
            # 提取basic_info中的粉丝数等关键数据
            if 'user' in basic_info:
                basic_user = basic_info['user']
                if 'user' in user_info:
                    # 如果粉丝数在basic_info中存在且有效，则使用它
                    if 'followers_count' in basic_user and basic_user['followers_count'] > 0:
                        user_info['user']['followers_count'] = basic_user['followers_count']
                    
                    # 同样处理好友数量
                    if 'friends_count' in basic_user and basic_user['friends_count'] > 0:
                        user_info['user']['friends_count'] = basic_user['friends_count']
            
            # 保留user_data变量，但主要使用user_info直接提取数据
            user_data = user_info.get('user', {})
            
            # 构建标准化的用户数据字典，补充更多详细信息
            row = {
                '用户主页的URL': f"https://weibo.com/u/{uid}",
                '唯一ID': uid,
                # 粉丝数：优先从user_data中获取，如无则尝试其他方法
                '用户粉丝数': user_data.get('followers_count', '') or extract_followers_count_from_all_sources(user_info),
                # 好友数：优先从user_data中获取，如无则尝试从friend_info提取
                '用户好友数': user_data.get('friends_count', '') or extract_friends_count(user_info),
                # 生日信息从detail_info获取
                '用户生日': user_info.get('birthday', ''),
                '用户性别': '女' if user_info.get('gender') == 'f' else '男' if user_info.get('gender') == 'm' else '未知',
                '用户所在地': user_info.get('location', '') or user_data.get('location', ''),
                '用户的认证状态': '已认证' if user_info.get('verified_url', '') or user_data.get('verified', False) else '未认证',
                '认证类型': determine_verified_type(user_info.get('verified_url', '')),
                '用户信用等级': user_info.get('sunshine_credit', {}).get('level', ''),
                '账号创建时间': user_info.get('created_at', ''),
                '个人描述': user_info.get('description', '') or user_data.get('description', ''),
                '实名认证': determine_real_auth_status(user_info),
                '公司': user_info.get('company', '') or user_info.get('career', {}).get('company', ''),
                '学校': user_info.get('education', {}).get('school', ''),
                '标签': ', '.join([label.get('name', '') for label in user_info.get('label_desc', []) if 'name' in label]) if user_info.get('label_desc') else '',
                '描述文本': user_info.get('desc_text', '')
            }
            
            return {'status': 'success', 'uid': uid, 'data': row}
            
        except Exception as e:
            error_message = f"获取用户 {uid} 信息失败: {str(e)}"
            print(f"\n[错误] {error_message}")
            return {'status': 'error', 'uid': uid, 'message': error_message, 'exception': str(e)}
    
    def worker(self, cookie_id, user_ids):
        """
        工作线程，处理分配给特定Cookie的用户ID列表
        
        Args:
            cookie_id: Cookie ID
            user_ids: 该Cookie负责处理的用户ID列表
        """
        total = len(user_ids)
        success_count = 0
        fail_count = 0
        consecutive_fails = 0  # 连续失败计数
        
        print(f"线程 {cookie_id} 开始处理 {total} 个用户")
        
        # 先验证cookie是否有效
        cookie_content = self.cookies_manager.get_cookie(cookie_id)
        if not cookie_content:
            print(f"\n[错误] 线程 {cookie_id} 的Cookie无效，无法启动")
            # 将该线程的所有任务添加到失败列表
            with self.fail_lock:
                for uid in user_ids:
                    if uid not in self.failed_ids:
                        self.failed_ids.append(uid)
                self.save_failed_ids()
            return
        
        # 进一步检查cookie有效性
        if not check_cookie_validity(cookie_content):
            print(f"\n[错误] 线程 {cookie_id} 的Cookie已失效，请更新")
            # 将该线程的所有任务添加到失败列表
            with self.fail_lock:
                for uid in user_ids:
                    if uid not in self.failed_ids:
                        self.failed_ids.append(uid)
                self.save_failed_ids()
            return
        
        # 处理每个用户ID
        for index, uid in enumerate(user_ids, 1):
            try:
                # 检查全局是否已处理
                if uid in self.processed_ids:
                    print(f"线程 {cookie_id} - 用户 {uid} 已被其他线程处理")
                    continue
                
                print(f"线程 {cookie_id} 处理 {index}/{len(user_ids)}: UID {uid}")
                
                # 处理用户
                result = self.process_user(uid, cookie_id)
                
                if result['status'] == 'success':
                    # 处理成功，重置连续失败计数
                    consecutive_fails = 0
                    
                    # 添加到结果列表
                    with self.results_lock:
                        self.results.append(result['data'])
                    
                    # 标记为已处理
                    with self.progress_lock:
                        self.processed_ids.add(uid)
                        # 每处理成功一个用户，就立即更新进度文件，确保断点续爬
                        save_progress(self.processed_ids)
                    
                    success_count += 1
                    
                    # 每5个用户保存一次结果
                    if success_count % 5 == 0:
                        self.save_results()
                    
                else:
                    # 处理失败
                    consecutive_fails += 1
                    fail_count += 1
                    
                    # 记录失败原因，避免重复爬取导致进度错误
                    print(f"\n[失败] 线程 {cookie_id} 处理用户 {uid} 失败: {result.get('message', '未知错误')}")
                    
                    with self.fail_lock:
                        if uid not in self.failed_ids:
                            self.failed_ids.append(uid)
                            self.save_failed_ids()
                    
                    # 如果连续失败超过5次，可能是cookie失效
                    if consecutive_fails >= 5:
                        print(f"\n[警告] 线程 {cookie_id} 连续失败 {consecutive_fails} 次，可能cookie已失效")
                        self.cookies_manager.mark_cookie_unavailable(cookie_id)
                        break
                    
                    # 失败时增加等待时间
                    time.sleep(random.uniform(3, 5))
                
            except Exception as e:
                print(f"\n[异常] 线程 {cookie_id} 处理用户 {uid} 时发生异常: {str(e)}")
                traceback.print_exc()
                
                consecutive_fails += 1
                fail_count += 1
                
                with self.fail_lock:
                    if uid not in self.failed_ids:
                        self.failed_ids.append(uid)
                        self.save_failed_ids()
                
                # 失败时增加等待时间
                time.sleep(random.uniform(3, 5))
                
                # 连续失败过多，退出
                if consecutive_fails >= 5:
                    print(f"\n[警告] 线程 {cookie_id} 连续失败 {consecutive_fails} 次，退出")
                    break
        
        # 线程结束时保存最终结果
        self.save_results()
        print(f"线程 {cookie_id} 已完成处理，成功: {success_count}，失败: {fail_count}")
    
    def run(self):
        """
        运行多线程爬取用户信息
        """
        # 拆分用户ID
        task_splits = self.split_user_ids()
        
        if not task_splits:
            print("没有任务需要处理或所有用户都已处理完成")
            return
        
        # 创建线程池
        with concurrent.futures.ThreadPoolExecutor(max_workers=len(task_splits)) as executor:
            # 提交任务
            futures = {executor.submit(self.worker, cookie_id, user_ids): cookie_id 
                      for cookie_id, user_ids in task_splits.items()}
            
            # 等待所有任务完成
            for future in concurrent.futures.as_completed(futures):
                cookie_id = futures[future]
                try:
                    future.result()  # 获取结果，如有异常会抛出
                except Exception as e:
                    print(f"\n[错误] 线程 {cookie_id} 执行失败: {str(e)}")
                    traceback.print_exc()
        
        # 所有线程完成后，保存最终结果
        self.save_results()
        self.save_failed_ids()
        
        # 输出统计信息
        success_count = len(self.processed_ids)
        failed_count = len(self.failed_ids)
        
        print("\n爬取完成！")
        print(f"成功获取: {success_count} 个用户信息")
        print(f"失败数量: {failed_count} 个用户")
        
        # 输出结果文件路径
        print(f"\n用户信息已保存到: {self.output_csv}")
        print(f"爬取进度已保存到: {os.path.join(self.progress_dir, 'user_crawl_progress.json')}")
        print(f"失败ID已保存到: {self.failed_ids_file}")


if __name__ == "__main__":
    try:
        # 确保必要的目录存在
        os.makedirs('Weibo', exist_ok=True)
        os.makedirs('Weibo/爬取进度', exist_ok=True)
        
        # 创建并运行多线程爬虫
        crawler = MultiThreadCrawler()
        crawler.run()
    except Exception as e:
        print(f"\n[错误] 程序执行出错: {str(e)}")
        traceback.print_exc() 